"""
数据库操作工具类
"""
import pymysql
from loguru import logger
from pymysql.cursors import DictCursor


class HandleMysql:
    def __init__(self, user, password, database, host, port=3306):
        self.conn = pymysql.connect(
            user=user,
            password=password,
            database=database,
            port=port,
            host=host,
            charset="utf8mb4",
            cursorclass=DictCursor  # 数据读取默认格式是元组，不方便确实数据字段，加上这个之后就会以字典格式存储数据。
        )
        self.cursor = self.conn.cursor()

    def query_data(self, sql, fetch_num=1, size=None):
        try:
            data = None
            result = self.cursor.execute(sql)
            logger.info(f"数据库的查询结果条数是：{result}")
            if result > 0:
                if fetch_num == 1:
                    data = self.cursor.fetchone()
                elif fetch_num >= 2:
                    data = self.cursor.fetchmany(size=size)
                elif fetch_num == -1:
                    data = self.cursor.fetchall()
            logger.info(f"数据库的查询数据结果为：{data}")
            return data
        except Exception as e:
            logger.error(f"数据库操作异常：{e}")
        finally:
            self.cursor.close()
            self.conn.close()


if __name__ == '__main__':
    from api_test_framework.config.settings import DB_CONFIG

    sql = 'select mobile_code,user_phone,rec_date from tz_sms_log where user_phone like "1785%";'
    print(HandleMysql(**DB_CONFIG).query_data(sql, 1))
